import pymysql
import threading
from contextlib import contextmanager

from pyassured.core.read_yaml import yaml_config
from pyassured.core.logger import Log

log = Log()


class MySQLConnectionPool:
    """MySQL连接池单例类"""
    _instance = None
    _lock = threading.Lock()
    _pools = {}  # 存储不同环境的连接池

    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
        return cls._instance

    def get_connection(self, env=None):
        """获取数据库连接"""
        pool_key = env or 'test'

        if pool_key not in self._pools:
            with self._lock:
                if pool_key not in self._pools:
                    self._pools[pool_key] = MySQLUtil(env=env)

        return self._pools[pool_key]

    def close_all_connections(self):
        """关闭所有连接池中的连接"""
        with self._lock:
            for pool in self._pools.values():
                try:
                    pool.close()
                except:
                    pass
            self._pools.clear()


class MySQLUtil():
    # 初始化
    def __init__(self, env=None):
        """
        初始化MySQL连接工具
        
        Args:
            env: 指定环境，如果不指定则使用配置文件中的默认环境
        """
        self.application_config = yaml_config.read_application_config("application.yaml")

        # 获取当前环境
        current_env = env or self.application_config.get('current_env', 'test')

        # 获取数据库配置
        dbconfig = self.application_config.get('dbconfig', {})

        # 支持多环境配置和向后兼容
        if isinstance(dbconfig, dict) and current_env in dbconfig:
            # 新的多环境配置格式
            env_config = dbconfig[current_env]
            log.info(f"使用环境 '{current_env}' 的数据库配置")
        elif isinstance(dbconfig, dict) and 'host' in dbconfig:
            # 旧的全局配置格式（向后兼容）
            env_config = dbconfig
            log.info("使用全局数据库配置（兼容模式）")
        else:
            # 配置错误或缺失
            raise ValueError(f"未找到环境 '{current_env}' 的数据库配置，请检查 application.yaml 文件")

        # 设置连接参数
        self.host = env_config.get('host')
        self.port = env_config.get('port')
        self.user = env_config.get('user')
        self.pwd = env_config.get('pwd')
        self.dbName = env_config.get('dbname') or env_config.get('dbName')  # 兼容两种命名

        # 验证必要参数
        if not all([self.host, self.port, self.user, self.pwd, self.dbName]):
            missing_params = [param for param, value in {
                'host': self.host, 'port': self.port, 'user': self.user,
                'pwd': self.pwd, 'dbName': self.dbName
            }.items() if not value]
            raise ValueError(f"数据库配置缺少必要参数: {', '.join(missing_params)}")

        log.info(f"数据库连接配置: {self.host}:{self.port}/{self.dbName} (用户: {self.user})")

    def connect(self):
        """建立数据库连接"""
        if not hasattr(self, 'db') or self.db is None:
            try:
                self.db = pymysql.connect(host=self.host, port=int(self.port), user=self.user, passwd=self.pwd,
                                          database=self.dbName, use_unicode=True, charset='utf8')
                self.cursor = self.db.cursor()
                log.info(f"建立新的数据库连接: {self.host}:{self.port}/{self.dbName}")
            except Exception as e:
                log.error(f"数据库连接失败: {e}")
                raise e
        else:
            # 检查连接是否有效
            try:
                self.db.ping(reconnect=True)
            except Exception as e:
                log.warning(f"连接已断开，重新连接: {e}")
                self._reconnect()

    def _reconnect(self):
        """重新连接数据库"""
        try:
            if hasattr(self, 'cursor') and self.cursor:
                self.cursor.close()
            if hasattr(self, 'db') and self.db:
                self.db.close()
        except:
            pass

        self.db = pymysql.connect(host=self.host, port=int(self.port), user=self.user, passwd=self.pwd,
                                  database=self.dbName, use_unicode=True, charset='utf8')
        self.cursor = self.db.cursor()
        log.info(f"重新建立数据库连接: {self.host}:{self.port}/{self.dbName}")

    @contextmanager
    def get_cursor(self):
        """获取数据库游标的上下文管理器"""
        self.connect()
        try:
            yield self.cursor
        except Exception as e:
            if hasattr(self, 'db') and self.db:
                self.db.rollback()
            raise e

    # 关闭数据库连接
    def close(self):
        """关闭数据库连接"""
        try:
            if hasattr(self, 'cursor') and self.cursor:
                self.cursor.close()
            if hasattr(self, 'db') and self.db:
                self.db.close()
            self.db = None
            self.cursor = None
            log.info("数据库连接已关闭")
        except Exception as e:
            log.error(f"关闭数据库连接时出错: {e}")

    # 查询操作，查询单条数据
    def get_one(self, sql):
        """查询单条数据，复用连接"""
        try:
            with self.get_cursor() as cursor:
                cursor.execute(sql)
                res = cursor.fetchone()
                return res
        except Exception as e:
            log.error(f"查询单条数据失败: {e}")
            return None

    # 查询操作，查询多条数据
    def get_all(self, sql):
        """查询多条数据，复用连接"""
        try:
            with self.get_cursor() as cursor:
                cursor.execute(sql)
                res = cursor.fetchall()
                return res
        except Exception as e:
            log.error(f"查询多条数据失败: {e}")
            return ()

    # 数据库插入、更新、删除操作
    def insert(self, sql):
        return self.__edit(sql)

    def update(self, sql):
        return self.__edit(sql)

    def delete(self, sql):
        return self.__edit(sql)

    def __edit(self, sql):
        """执行编辑操作，复用连接"""
        try:
            with self.get_cursor() as cursor:
                count = cursor.execute(sql)
                self.db.commit()
                return count
        except Exception as e:
            log.error(f"执行编辑操作失败: {e}")
            return 0


# 提供便捷的连接池访问函数
def get_db_connection(env=None):
    """获取数据库连接的便捷函数"""
    pool = MySQLConnectionPool()
    return pool.get_connection(env=env)


def close_all_db_connections():
    """关闭所有数据库连接的便捷函数"""
    pool = MySQLConnectionPool()
    pool.close_all_connections()


if __name__ == '__main__':
    mysql = MySQLUtil()
